import matplotlib.pyplot as plt
import matplotlib.animation as animation


# 冒泡排序算法
def bubble_sort(arr):
    n = len(arr)
    steps = []
    comparisons = []
    steps.append(list(arr))  # 初始状态
    comparisons.append((-1, -1))  # 初始状态没有比较
    for i in range(n):
        for j in range(0, n - i - 1):
            comparisons.append((j, j + 1))  # 记录比较的索引
            if arr[j] > arr[j + 1]:
                arr[j], arr[j + 1] = arr[j + 1], arr[j]
                steps.append(list(arr))  # 每次交换后记录状态
            else:
                steps.append(list(arr))  # 即使没有交换也记录状态
    return steps, comparisons


# 初始化数据
n = 10
array = [3, 6, 2, 8, 4, 1, 7, 5, 9, 0]

# 获取排序过程中的所有状态和比较情况
steps, comparisons = bubble_sort(array)

# 绘制动画
fig, ax = plt.subplots()
bar_rects = ax.bar(range(len(array)), steps[0], align="edge")
ax.set_xlim(0, n)
ax.set_ylim(0, int(max(array) * 1.1))
ax.set_title("Bubble Sort Visualization")
ax.set_xlabel("Index")
ax.set_ylabel("Value")


# 更新每一帧
def update_plot(step_index):
    step = steps[step_index]
    comp = comparisons[step_index]
    for rect, val in zip(bar_rects, step):
        rect.set_height(val)
        rect.set_color('b')
    if comp[0] != -1 and comp[1] != -1:
        bar_rects[comp[0]].set_color('r')
        bar_rects[comp[1]].set_color('r')
    return bar_rects


# 创建动画
anim = animation.FuncAnimation(fig, update_plot, frames=range(len(steps)), interval=1000, repeat=False)

plt.show()
